from TCMLLMHandler import TCMDBHandler


class TCMLLM:
    def __init__(self, db_path: str):
        self.db = TCMDBHandler(db_path)
        # 新增：别名联想规则（解决用户用别名提问的问题）
        self.alias_map = self._load_alias_map()  # 从数据库加载别名映射表

    def _load_alias_map(self) -> Dict[str, str]:
        """从数据库预加载别名→正式名称的映射（提升查询成功率）"""
        sql = "SELECT herb_name, alias FROM tcm_herb WHERE alias IS NOT NULL;"
        alias_map = {}
        for row in self.db.connection.execute(sql).fetchall():
            for alias in row["alias"].split(";"):
                alias_map[alias.strip()] = row["herb_name"]
        return alias_map

    def _parse_user_query(self, query: str):
        """增强：将别名转换为正式药材名称"""
        query = self.alias_map.get(query, query)  # 替换别名→正式名
 